#!/usr/bin/env python
"""
This script creates two plots showing the performance comparison between inputEmployeeHead() and inputEmployeeTail()
functions, with data size on the x-axis and execution time on the y-axis.
"""

import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit

# Test data from your output
# Data size (employees) vs execution time (ms)
# Format: [data_size, head_time_ms, tail_time_ms]
test_data = [
    [1000, 0, 12],
    [2000, 1, 87],
    [3000, 1, 215],
    [4000, 2, 532],
    [5000, 3, 1004],
    [6000, 3, 2237],
    [7000, 3, 3811],
    [8000, 4, 6544],
    [9000, 4, 9664],
    [10000, 7, 16619],
]

# Extract data for plotting
sizes = [row[0] for row in test_data]
head_times = [row[1] for row in test_data]
tail_times = [row[2] for row in test_data]

# Define fitting functions
def linear_func(x, a, b):
    """
    Linear function: y = ax + b
    """
    return a * x + b

def quadratic_func(x, a, b):
    """
    Quadratic function: y = ax^2 + bx
    """
    return a * x**2 + b * x

# Create figure with two subplots
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

# Plot 1: inputEmployeeHead() - O(n) behavior
# Fit linear model to head times
popt_head, pcov_head = curve_fit(linear_func, sizes, head_times)
head_fitted = linear_func(np.array(sizes), *popt_head)

# Plot original data points
ax1.scatter(sizes, head_times, color='blue', label='Original Data Points', s=50)

# Plot fitted curve
ax1.plot(sizes, head_fitted, color='red', linewidth=2, label=f'Fitted Curve (y={popt_head[0]:.3f}x+{popt_head[1]:.3f})')

# Add labels and title
ax1.set_xlabel('Data Size (Number of Employees)', fontsize=12)
ax1.set_ylabel('Execution Time (ms)', fontsize=12)
ax1.set_title('inputEmployeeHead() Performance', fontsize=14, fontweight='bold')
ax1.legend(fontsize=10)
ax1.grid(True, alpha=0.3)

# Plot 2: inputEmployeeTail() - O(n²) behavior
# Fit quadratic model to tail times
popt_tail, pcov_tail = curve_fit(quadratic_func, sizes, tail_times)
tail_fitted = quadratic_func(np.array(sizes), *popt_tail)

# Plot original data points
ax2.scatter(sizes, tail_times, color='blue', label='Original Data Points', s=50)

# Plot fitted curve
ax2.plot(sizes, tail_fitted, color='red', linewidth=2, label=f'Fitted Curve (y={popt_tail[0]:.3e}x²+{popt_tail[1]:.3f}x)')

# Add labels and title
ax2.set_xlabel('Data Size (Number of Employees)', fontsize=12)
ax2.set_ylabel('Execution Time (ms)', fontsize=12)
ax2.set_title('inputEmployeeTail() Performance', fontsize=14, fontweight='bold')
ax2.legend(fontsize=10)
ax2.grid(True, alpha=0.3)

# Adjust layout and display
plt.tight_layout()
plt.savefig('performance_comparison.png', dpi=300, bbox_inches='tight')
print("Performance comparison plots saved as 'performance_comparison.png'")

# Show the plots
plt.show()